#include "./includes/utility.h"


int main (int argc, const char * argv[]) {

    //shared key
    int KEYLEN;
    KEYLEN = EVP_MAX_KEY_LENGTH;
    unsigned char sharedkey[KEYLEN];
    //Buffers
    unsigned char *plaintext,*ciphertext,*datatosend,*cipherkey;
    int pt_len,ct_len,ck_len,h_len,datasize,enc_bytes,dec_bytes,bsize;
    char* err;
    //Directory
    char* client_dir = "./received/";
    char* filename;
    //Server public key
    FILE* pubkey;
    //RSA data
    RSA* rsa;
    char * file_pem_pub;
    file_pem_pub = "pub.pem";
    //socket data
    struct sockaddr_in sst;
    int socket;
    int srv_sock,udp_sock,fdmax,bytes_rcv;
    int quit=1;
    int retSelect=-1;
    struct sockaddr_in srv_addr;
    char str_console='>';
    int console_flag=1;
    char buf[BUFLEN];
    char username[USERLEN];
    char password[PSSWLEN];
    char client_nonce[NONCELEN];
    char server_nonce[NONCELEN];
    //Select data
    fd_set read_fds;
    fd_set master;

    if (argc != 3)
	    error("Parameters Error: Server IP and PORT needed");
    while(1){
		//Init TCP connection with server
    	//1st protocol step [(M1) : C -> S]
		TCP_init(&srv_sock,&srv_addr,argv[1],atoi(argv[2]));
		TCP_connect(srv_sock,&srv_addr);
		printf("Connected to Server %s (Port %s)\n",argv[1],argv[2]);
		//printf("[(M1) : C -> S]\n");
		//RSA init
		OpenSSL_add_all_algorithms();
		rsa = RSA_new();
		//Retrieve server's public key
		pubkey = fopen(file_pem_pub,"r");
		PEM_read_RSA_PUBKEY(pubkey,&rsa,NULL,NULL);
		fclose(pubkey);
		bsize = RSA_size(rsa);
		//2st protocol message [(M2) : S -> C]
		TCP_recv(srv_sock,(void*)&ct_len,sizeof(int));
		ciphertext = (unsigned char*)malloc(ct_len);
		TCP_recv(srv_sock,(void*)ciphertext,ct_len);
		plaintext = (unsigned char*)malloc(bsize);
		//printf("[(M2) : S -> C] : ");
		fflush(stdout);
		dec_bytes = RSA_public_decrypt(ct_len,ciphertext,plaintext,rsa,RSA_PKCS1_PADDING);
		if(dec_bytes == -1) {
			err = (char*)malloc(130);
			ERR_load_CRYPTO_strings();
			ERR_error_string(ERR_get_error(),err);
			printf("%s\n",err);
			perror("Error: ");
			free(err);
			free(ciphertext);
			free(plaintext);
			RSA_free(rsa);
			TCP_close(srv_sock);
			continue;
		}
		if(strncmp(plaintext,"!INIT",5) != 0){
			printf("Init message not valid! Connection will be closed\n");
			free(ciphertext);
			free(plaintext);
			RSA_free(rsa);
			TCP_close(srv_sock);
			continue;
		}
		//copy server_nonce
		memcpy(server_nonce,&plaintext[5],NONCELEN);
		free(ciphertext);
		free(plaintext);
		RSA_free(rsa);
		printf("Init message OK\n");
		//-----------------------------------------
		//---------Authentication phase------------
		//-----------------------------------------
		//3rd protocol step [(M3) : C -> S]
		printf("username: ");
		memset(username,0,USERLEN);
		sgets(username,USERLEN);
		printf("password: ");
		memset(password,0,PSSWLEN);
		sgets(password,PSSWLEN);
		//Create sharedkey
		RAND_bytes(sharedkey,KEYLEN);
		//RSA init
		OpenSSL_add_all_algorithms();
		rsa = RSA_new();
		//Retrieve server's public key
		pubkey = fopen(file_pem_pub,"r");
		PEM_read_RSA_PUBKEY(pubkey,&rsa,NULL,NULL);
		fclose(pubkey);
		//(Sharedkey|server_nonce|password) encryption with server's public key
		bsize = RSA_size(rsa);
		ck_len = bsize;
		cipherkey = (char*)malloc(ck_len);
		pt_len = KEYLEN+NONCELEN+PSSWLEN;
		plaintext = (unsigned char*)malloc(pt_len);
		memcpy(plaintext,sharedkey,KEYLEN);
		memcpy(&plaintext[KEYLEN],server_nonce,NONCELEN);
		memcpy(&plaintext[KEYLEN+NONCELEN],password,PSSWLEN);
		enc_bytes = RSA_public_encrypt(pt_len,plaintext,cipherkey,rsa,RSA_PKCS1_PADDING);
		if(enc_bytes == -1 || enc_bytes < bsize) {
		  err = (char*)malloc(130);
		  ERR_load_CRYPTO_strings();
		  ERR_error_string(ERR_get_error(),err);
		  printf("%s\n",err);
		  perror("Error: ");
		  free(err);
		  exit(1);
		}
		//Nonce generation (8 bytes = 64 bits)
		create_nonce((unsigned char*) client_nonce);
		//concatenation
		pt_len = USERLEN+PSSWLEN+NONCELEN;
		plaintext = (char*)malloc(pt_len);
		memcpy(plaintext,username,USERLEN);
		memcpy(&plaintext[USERLEN],password,PSSWLEN);
		memcpy(&plaintext[USERLEN+PSSWLEN],client_nonce,NONCELEN);
		//Ciphertext init
		bsize = block_size();
		ct_len = pt_len + bsize;
		ciphertext = (unsigned char*)malloc(ct_len);
		//Encryption with shared key of (username|h(password)|nonce)
		ct_len = encrypt(sharedkey,KEYLEN,ciphertext,pt_len,plaintext);
		//Data to send init
		datasize = ck_len+ct_len;
		datatosend = (unsigned char*)malloc(datasize);
		//copy cipherkey into datatosend
		memcpy(datatosend,cipherkey,ck_len);
		//concatenation of ciphertext to cipherkey
		memcpy(&datatosend[ck_len],ciphertext,ct_len);
		//send data length
		TCP_send(srv_sock,(void*)&datasize,sizeof(datasize));
		//send data (cipherkey|ciphertext)
		TCP_send(srv_sock,(void*)datatosend,datasize);
		//Free resources
		free(datatosend);
		free(plaintext);
		free(ciphertext);
		//printf("[(M3) : C -> S]\n");
		//4th protocol step [(M4) : S -> C]
		//wait server response
		char* msg;
		int nb;
		msg = TCP_recv_enc_msg(sharedkey,KEYLEN,srv_sock,&nb);
		//printf("[(M4) : S -> C] : ");
		if(strcmp(msg,"!ER00") == 0){
			printf("Nonce not correct!\nConnection closed\n\n");
			TCP_close(srv_sock);
			free(msg);
			continue;
		}
		else if(strcmp(msg,"!ER01") == 0){
			printf("Username or Password not correct!\nConnection closed\n\n");
			TCP_close(srv_sock);
			free(msg);
			continue;
		}
		else if(strcmp(msg,"!ER02") == 0){
			printf("Password not correct!\nConnection closed\n\n");
			TCP_close(srv_sock);
			free(msg);
			continue;
		}
		else if(strcmp(msg,"!ER03") == 0){
			printf("Client already connected\nConnection closed\n\n");
			TCP_close(srv_sock);
			free(msg);
			continue;
		}
		else{
			char rcv_cmd[6];
			char rcv_nonce[NONCELEN];
			memcpy(rcv_cmd,msg,5);
			memcpy(rcv_nonce,&msg[5],NONCELEN);
			rcv_cmd[5] = '\0';
			//check message received
			if(strcmp(rcv_cmd,"!OK01") == 0) {
				if((memcmp(rcv_nonce,client_nonce,NONCELEN) != 0)){
					printf("Nonce not valid!\nConnection closed\n\n");
					TCP_close(srv_sock);
					free(msg);
					continue;
				}
				if((memcmp(&msg[5+NONCELEN],sharedkey,KEYLEN) != 0)){
					printf("Sharedkey not valid!\nConnection closed\n\n");
					TCP_close(srv_sock);
					free(msg);
					continue;
				}
				printf("Authentication OK\n");
			}
			else{
				printf("Received command not valid!\nConnection closed\n\n");
				TCP_close(srv_sock);
				free(msg);
				continue;
			}
			free(msg);
			break;
		}
	}
	//Select FD
	FD_ZERO(&master);
	FD_SET(STDIN_FILENO,&master);
	FD_SET(srv_sock,&master);
	fdmax=srv_sock;

	while(quit){
		int i;
		read_fds=master;
		if(console_flag)
			printf("%c ",str_console);
		fflush(stdout);
		retSelect=select(fdmax+1, &read_fds, NULL, NULL, NULL);
		if (retSelect == -1) {
			error("Select error");
		}
		for (i=0; i<=fdmax; i++) {
			if (FD_ISSET(i,&read_fds)) {
				//Bytes available in stdin
				if (i == STDIN_FILENO) {
					//read from stdin
					sgets(buf, BUFLEN);
					//check for commands
					char cmd[5];
					memcpy(cmd,buf,5);
					console_flag = 0;
					if (strcmp(cmd, "!QUIT")==0) {
						TCP_send_enc_msg(sharedkey,KEYLEN,srv_sock,buf,strlen(buf));
						FD_CLR(srv_sock, &master);
						TCP_close(srv_sock);
						quit=0;
						break;
					} else if (strcmp(cmd, "!HELP") == 0) {
						TCP_send_enc_msg(sharedkey,KEYLEN,srv_sock,buf,strlen(buf));
					} else if (strcmp(cmd, "!LIST") == 0) {
						TCP_send_enc_msg(sharedkey,KEYLEN,srv_sock,buf,strlen(buf));
					} else if (strcmp(cmd, "!UPLD") == 0) {
						FILE *fp;
						filename = str_trim(&buf[5]);
						fp = fopen(filename,"r");
						if(fp == NULL){
							printf("Can't open %s\n",filename);
							console_flag = 1;
						}
						else{
							TCP_send_enc_msg(sharedkey,KEYLEN,srv_sock,buf,strlen(buf));
							fclose(fp);
						}
					} else if (strcmp(cmd, "!DWLD") == 0) {
						TCP_send_enc_msg(sharedkey,KEYLEN,srv_sock,buf,strlen(buf));
						filename = str_trim(&buf[5]);
					}  else {
						printf("wrong command\n");
						console_flag = 1;
					}

				}
				//message from server
				if (i == srv_sock) {
					char* msg;
					int nb;
					fflush(stdout);
					msg = TCP_recv_enc_msg(sharedkey,KEYLEN,i,&nb);
					if(msg == NULL){
						quit=0;
						break;
					}
					else if(strcmp(msg,"!DWLD") == 0){
						char* fpath;
						fpath = (char*)malloc(strlen(filename)+strlen(client_dir));
						strcpy(fpath,client_dir);
						strcat(fpath,filename);
						printf("download in progress...\n");
						receive_enc_file(fpath, sharedkey,KEYLEN,srv_sock);
						printf("download completed\n");
						free(fpath);
						free(filename);
					}
					else if(strcmp(msg,"!UPLD") == 0){
						printf("upload in progress...\n");
						send_enc_file(filename, sharedkey,EVP_MAX_KEY_LENGTH,srv_sock);
						printf("file sent\n");
						free(filename);
					}
					else{
						printf("%s\n",msg);
					}
					free(msg);
					console_flag = 1;
				}
			}
		}
	}
	printf("Bye Bye\n");
    return 0;
}


